import subprocess
from pycparser import c_parser, c_ast
import json

from utils.parseData import response_json


# --- 1. 预处理 C 代码 ---
def preprocess_code(filename):
    result = subprocess.run(['gcc', '-E', filename], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, text=True)
    return result.stdout

# --- 2. 将 AST 节点还原为 C 代码字符串 ---
def reconstruct_code(node):
    if isinstance(node, c_ast.Assignment):
        return f"{reconstruct_code(node.lvalue)} {node.op} {reconstruct_code(node.rvalue)}"
    elif isinstance(node, c_ast.ID):
        return node.name
    elif isinstance(node, c_ast.Constant):
        return node.value
    elif isinstance(node, c_ast.BinaryOp):
        return f"({reconstruct_code(node.left)} {node.op} {reconstruct_code(node.right)})"
    elif isinstance(node, c_ast.UnaryOp):
        return f"{node.op}{reconstruct_code(node.expr)}"
    elif isinstance(node, c_ast.FuncCall):
        name = reconstruct_code(node.name)
        args = ', '.join([reconstruct_code(a) for a in node.args.exprs]) if node.args else ''
        return f"{name}({args})"
    elif isinstance(node, c_ast.Cast):
        return f"({reconstruct_code(node.to_type)}){reconstruct_code(node.expr)}"
    elif isinstance(node, c_ast.TypeDecl):
        return node.declname
    elif isinstance(node, c_ast.Decl):
        return node.name
    elif isinstance(node, c_ast.ArrayRef):
        return f"{reconstruct_code(node.name)}[{reconstruct_code(node.subscript)}]"
    elif isinstance(node, c_ast.Compound):
        return '{...}'
    elif isinstance(node, c_ast.Return):
        return f"return {reconstruct_code(node.expr)}"
    elif isinstance(node, c_ast.If):
        return f"if ({reconstruct_code(node.cond)})"
    elif isinstance(node, c_ast.While):
        return f"while ({reconstruct_code(node.cond)})"
    return type(node).__name__

# --- 3. 控制流图构建类 ---
class CFGBuilder(c_ast.NodeVisitor):
    def __init__(self):
        self.nodes = []  # Store nodes in a list
        self.edges = []  # Store edges in a list
        self.count = 1  # Start node ID from 1
        self.prev = []

    def add_node(self, label):
        # 为每个节点添加唯一编号
        node_id = self.count
        self.nodes.append({"id": node_id, "label": label})
        self.count += 1

        # 为当前节点添加边
        if self.prev:
            for p in self.prev:
                self.edges.append({"source": p, "target": node_id})

        self.prev = [node_id]
        return node_id

    def visit_Compound(self, node):
        for stmt in node.block_items or []:
            self.visit(stmt)

    def visit_FuncDef(self, node):
        # Add start node for the function
        start_node_id = self.add_node(f"Function: {node.decl.name}")
        self.prev = [start_node_id]
        self.visit(node.body)

    def visit_If(self, node):
        cond = reconstruct_code(node.cond)
        cond_id = self.add_node(f"If: {cond}")

        prev_copy = self.prev[:]
        self.prev = [cond_id]
        self.visit(node.iftrue)
        true_exit = self.prev[:]

        self.prev = [cond_id]
        if node.iffalse:
            self.visit(node.iffalse)
            false_exit = self.prev[:]
        else:
            false_exit = [cond_id]

        self.prev = true_exit + false_exit

    def visit_While(self, node):
        cond = reconstruct_code(node.cond)
        cond_id = self.add_node(f"While: {cond}")
        loop_start = cond_id

        self.prev = [cond_id]
        self.visit(node.stmt)
        if self.prev:
            for p in self.prev:
                self.edges.append({"source": p, "target": loop_start})
        self.prev = [cond_id]

    def visit_Assignment(self, node):
        self.add_node(reconstruct_code(node))

    def visit_Decl(self, node):
        if node.init:
            self.add_node(reconstruct_code(node))

    def visit_FuncCall(self, node):
        self.add_node(f"Call: {reconstruct_code(node)}")

    def visit_Return(self, node):
        self.add_node(reconstruct_code(node))

def generate_cfg_json(c_code):
    parser = c_parser.CParser()
    try:
        ast = parser.parse(c_code)
    except Exception as e:
        return response_json(515, "Error", str(e))


    builder = CFGBuilder()
    builder.visit(ast)
    demo_json = {
        "nodes": builder.nodes,
        "edges": builder.edges
    }
    print(demo_json)


    # DEMOJSON = {
    #   "nodes": [
    #     { "id": 1, "label": "Start" },
    #     { "id": 2, "label": "Decl prime" },
    #     { "id": 3, "label": "Decl main" },
    #     { "id": 4, "label": "Decl i" },
    #     { "id": 5, "label": "Decl j" },
    #     { "id": 6, "label": "Decl v" },
    #     { "id": 7, "label": "Assignment prime[0] = 2" },
    #     { "id": 8, "label": "FuncCall printf (\" %lld\")" },
    #     { "id": 9, "label": "Assignment i = 1" },
    #     { "id": 10, "label": "Assignment v = 3" },
    #     { "id": 11, "label": "While i < 400" },
    #     { "id": 12, "label": "Assignment j = 0" },
    #     { "id": 13, "label": "While j < i" },
    #     { "id": 14, "label": "If prime[j] * prime[j] > v" },
    #     { "id": 15, "label": "Assignment j = i - 1" },
    #     { "id": 16, "label": "If v % prime[j] == 0" },
    #     { "id": 17, "label": "Assignment j = i" },
    #     { "id": 18, "label": "Assignment j = j + 1" },
    #     { "id": 19, "label": "Assignment prime[i] = v" },
    #     { "id": 20, "label": "FuncCall printf (\" %lld\")" },
    #     { "id": 21, "label": "Assignment i = i + 1" },
    #     { "id": 22, "label": "Assignment v = v + 2" },
    #     { "id": 23, "label": "FuncCall printf (\"\\n\")" },
    #     { "id": 24, "label": "End" }
    #   ],
    #   "edges": [
    #     { "source": 1, "target": 2 },
    #     { "source": 2, "target": 3 },
    #     { "source": 3, "target": 4 },
    #     { "source": 4, "target": 5 },
    #     { "source": 5, "target": 6 },
    #     { "source": 6, "target": 7 },
    #     { "source": 7, "target": 8 },
    #     { "source": 8, "target": 9 },
    #     { "source": 9, "target": 10 },
    #     { "source": 10, "target": 11 },
    #     { "source": 11, "target": 12 },
    #     { "source": 12, "target": 13 },
    #     { "source": 13, "target": 14 },
    #     { "source": 14, "target": 15 },
    #     { "source": 14, "target": 16 },
    #     { "source": 15, "target": 19 },
    #     { "source": 16, "target": 17 },
    #     { "source": 16, "target": 18 },
    #     { "source": 17, "target": 19 },
    #     { "source": 18, "target": 19 },
    #     { "source": 19, "target": 20 },
    #     { "source": 20, "target": 21 },
    #     { "source": 21, "target": 22 },
    #     { "source": 22, "target": 11 },
    #     { "source": 11, "target": 23 },
    #     { "source": 23, "target": 24 }
    #   ]
    # }
    # print(builder.to_json());
    return response_json(200, "Success", json.dumps(demo_json));

# # 测试
# c_code = """
# int main() {
#     int a = 5;
#     if (a > 3) {
#         a = a - 1;
#     } else {
#         a = a + 1;
#     }
#     return a;
# }
# """
#
# cfg_json = generate_cfg_json(c_code)
# if cfg_json:
#     print(json.dumps(cfg_json, indent=4))
